import glob
import math
import os
import wave

import pywt
import scipy.io.wavfile as wav
import numpy as np

from scipy import signal
import matplotlib.pyplot as plt
from scipy.signal import wiener

from fun_wav import GaussSmooth
from nextpow2 import nextpow2

count = 0

def CutBi(sig, length):
    sigLen = len(sig)
    sig = sig[length:sigLen-length]
    return sig

def LowFiltWav(signalWav, N, lowPass):
    b, a = signal.butter(N=N, Wn=lowPass, btype='lowpass')
    filtSignal = signal.filtfilt(b, a, signalWav)  # signalWav为要过滤的信号
    return filtSignal

def HighFiltWav(signalWav, N, highPass):
    b, a = signal.butter(N=N, Wn=highPass, btype='highpass')
    filtSignal = signal.filtfilt(b, a, signalWav)  # signalWav为要过滤的信号
    return filtSignal

def SaveWav(waveData, fileName):
    # cast to the type of short
    channels = 1
    sampWidth = 2
    frameRate = 9600

    maxValue = np.max(abs(waveData)) / 0.8
    waveData = waveData / maxValue
    waveData = waveData * 32765

    plt.plot(waveData)
    plt.title('res after')
    plt.show()

    waveData = waveData.astype(np.short)
    # open a wav document
    f = wave.open(fileName, "wb")
    # set wav params
    f.setnchannels(channels)
    f.setsampwidth(sampWidth)
    f.setframerate(frameRate)
    # turn the data to string
    f.writeframes(waveData.tostring())
    f.close()

def GetSpeed(angle, fc, fs):
    '''
    计算速度分为四步
    第一步：让angle乘以lamda/(2*pi)得到测得的距离
    第二步：相邻距离之间相减，得到距离差
    第三步：距离差除以10ms，得到速度
    '''
    interval = 0.01
    interPoint = int(fs * interval)
    angleLen = len(angle)
    batch = angleLen // interPoint
    wavLen = 340 / fc
    # 第一步
    angle = angle * wavLen / (2 * np.pi)
    speedRes = []
    for i in range(batch-1):
        start = i * interPoint
        end = (i+1) * interPoint
        speed = (angle[end] - angle[start]) / interval
        speedRes.append(speed)
    speedRes = np.array(speedRes)

    # speedRes = DirectMinus(speedRes, fs)
    # speedRes = wiener(speedRes)
    # interval = 0.01
    # interPoint = int(fs * interval)
    # angleLen = len(angle)
    # realLen = angleLen // interPoint
    # index = np.linspace(0, realLen, realLen, endpoint=False) * interPoint
    # index = index.astype(np.int)

    return speedRes

def GetCw(type, time, start, end, interval):
    fcList = np.arange(start, end+1, interval)
    cwList = []
    if type == 'cos':
        for fc in fcList:
            cw = np.cos(2 * np.pi * fc * time)
            cwList.append(cw)
    else:
        for fc in fcList:
            cw = -np.sin(2 * np.pi * fc * time)
            cwList.append(cw)
    return cwList

def WaveletTrans(signalWav, threshold):
    wavStr = 'sym4'
    # wavStr = 'db8'
    w = pywt.Wavelet(wavStr)  # 选用Daubechies8小波

    maxlev = pywt.dwt_max_level(len(signalWav), w.dec_len)
    print("maximum level is " + str(maxlev))
    threshold = threshold  # Threshold for filtering
    # Decompose into wavelet components, to the level selected:
    coeffs = pywt.wavedec(signalWav, wavStr, level=maxlev)  # 将信号进行小波分解
    for i in range(1, len(coeffs)):
        coeffs[i] = pywt.threshold(coeffs[i], threshold * max(coeffs[i]), mode='hard')  # 将噪声滤波
    datarec = pywt.waverec(coeffs, wavStr)
    return datarec

def CalSingalDop(signalWav, cwcos, cwsin, lowTh, fc, fs):

    real = signalWav * cwcos
    image = signalWav * cwsin
    lowPass = lowTh * 2/ fs
    th = 0.05

    real = LowFiltWav(real, 3, lowPass)
    image = LowFiltWav(image, 3, lowPass)

    divePer = 3
    real = CutBi(real, int(fs/divePer))
    image = CutBi(image, int(fs/divePer))

    # plt.plot(real)
    # plt.show()

    # wavTime = np.linspace(0, len(real) / fs, len(real), endpoint=False)
    # plt.plot(wavTime, real)
    # plt.xlabel('time')
    # plt.title(fc)
    # plt.show()


    image = image * (1j)
    mix = real + image
    # mix = CutBi(mix, fs//3)
    angle = np.angle(mix)
    angle = np.unwrap(angle)

    # plt.plot(angle)
    # plt.title(f'angle {fc}')
    # plt.show()

    speedRes = GetSpeed(angle, fc, fs)
    return speedRes

def DirectMinus(signalWav, fs):
    x = signalWav
    # 计算参数
    len_ = 20 * fs // 1000  # 样本中帧的大小
    PERC = 50  # 窗口重叠占帧的百分比
    len1 = len_ * PERC // 100  # 重叠窗口
    len2 = len_ - len1  # 非重叠窗口
    # 设置默认参数
    Thres = 3
    Expnt = 2.0
    beta = 0.002
    G = 0.9
    # 初始化汉明窗
    win = np.hamming(len_)
    # normalization gain for overlap+add with 50% overlap
    winGain = len2 / sum(win)

    # Noise magnitude calculations - assuming that the first 5 frames is noise/silence
    nFFT = 2 * 2 ** (nextpow2(len_))

    noise_mean = np.zeros(nFFT)

    j = 0
    for k in range(1, 6):
        noise_mean = noise_mean + abs(np.fft.fft(win * x[j:j + len_], nFFT))
        j = j + len_
    noise_mu = noise_mean / 5

    # --- allocate memory and initialize various variables
    k = 1
    img = 1j
    x_old = np.zeros(len1)
    Nframes = len(x) // len2 - 1
    xfinal = np.zeros(Nframes * len2)

    # =========================    Start Processing   ===============================
    for n in range(0, Nframes):
        # Windowing
        insign = win * x[k - 1:k + len_ - 1]
        # compute fourier transform of a frame
        spec = np.fft.fft(insign, nFFT)
        # compute the magnitude
        sig = abs(spec)

        # save the noisy phase information
        theta = np.angle(spec)
        SNRseg = 10 * np.log10(np.linalg.norm(sig, 2) ** 2 / np.linalg.norm(noise_mu, 2) ** 2)

        def berouti(SNR):
            if -5.0 <= SNR <= 20.0:
                a = 4 - SNR * 3 / 20
            else:
                if SNR < -5.0:
                    a = 5
                if SNR > 20:
                    a = 1
            return a

        def berouti1(SNR):
            if -5.0 <= SNR <= 20.0:
                a = 3 - SNR * 2 / 20
            else:
                if SNR < -5.0:
                    a = 4
                if SNR > 20:
                    a = 1
            return a

        if Expnt == 1.0:  # 幅度谱
            alpha = berouti1(SNRseg)
        else:  # 功率谱
            alpha = berouti(SNRseg)
        #############
        sub_speech = sig ** Expnt - alpha * noise_mu ** Expnt
        # 当纯净信号小于噪声信号的功率时
        diffw = sub_speech - beta * noise_mu ** Expnt

        # beta negative components

        def find_index(x_list):
            index_list = []
            for i in range(len(x_list)):
                if x_list[i] < 0:
                    index_list.append(i)
            return index_list

        z = find_index(diffw)
        if len(z) > 0:
            # 用估计出来的噪声信号表示下限值
            sub_speech[z] = beta * noise_mu[z] ** Expnt
            # --- implement a simple VAD detector --------------
        if SNRseg < Thres:  # Update noise spectrum
            noise_temp = G * noise_mu ** Expnt + (1 - G) * sig ** Expnt  # 平滑处理噪声功率谱
            noise_mu = noise_temp ** (1 / Expnt)  # 新的噪声幅度谱
        # flipud函数实现矩阵的上下翻转，是以矩阵的“水平中线”为对称轴
        # 交换上下对称元素
        sub_speech[nFFT // 2 + 1:nFFT] = np.flipud(sub_speech[1:nFFT // 2])
        x_phase = (sub_speech ** (1 / Expnt)) * (
                    np.array([math.cos(x) for x in theta]) + img * (np.array([math.sin(x) for x in theta])))
        # take the IFFT

        xi = np.fft.ifft(x_phase).real
        # --- Overlap and add ---------------
        xfinal[k - 1:k + len2 - 1] = x_old + xi[0:len1]
        x_old = xi[0 + len1:len_]
        k = k + len2
    return xfinal * winGain


def SingleFileDop(path, count=None):
    print(path)
    fs, signalWav = wav.read(path)
    signalWav = signalWav

    wavTime = np.linspace(0, len(signalWav)/fs, len(signalWav), endpoint=False)
    # plt.plot(wavTime, signalWav)
    # plt.xlabel('time')
    # plt.title('original')
    # plt.show()

    fstart = 15550
    fend = 18350
    finter = 350
    fcList = np.arange(fstart, fend + 1, finter)
    signalLen = len(signalWav)
    time = np.linspace(0, signalLen / fs, signalLen, endpoint=False)

    cwcosList = GetCw('cos', time, fstart, fend, finter)
    cwsinList = GetCw('sin', time, fstart, fend, finter)

    audioTh = 14000
    lowTh = 100
    audioPass = 2 * audioTh / fs
    signalWav = HighFiltWav(signalWav, 3, audioPass)

    freqWavNum = len(fcList)
    speedAll = []
    weights = np.array([1] * freqWavNum)
    for i in range(len(cwsinList)):
        speed = CalSingalDop(signalWav=signalWav,
                             cwcos=cwcosList[i],
                             cwsin=cwsinList[i],
                             lowTh=lowTh,
                             fc=fcList[i],
                             fs=fs)
        speedAll.append(speed)

    speedAll = np.array(speedAll)
    resNum = speedAll.shape[0]
    firstRes = speedAll[0]
    time = time[0:len(firstRes)]
    timeShow = np.arange(0, len(firstRes) / 100, 0.01)

    # plt.plot(firstRes)
    # plt.title('res before')
    # plt.show()
    # SaveWav(firstRes, 'have_3.wav')

    # sinSec = 44100
    # allVar = np.zeros((len(speedAll), ))
    # for i in range(len(speedAll)):
    #     allVar[i] = np.var(speedAll[i][:sinSec])
    # varSum = np.sum(allVar)
    # weights = allVar / varSum

    speedRes = np.zeros((len(firstRes), ))
    for i in range(resNum):
        speedRes = np.sum(speedAll, axis=0) / freqWavNum
        speedRes += speedAll[i] * weights[i]

        curSig = speedAll[i]
        plt.plot(curSig)
        plt.title(f'Freq is {fcList[i]}')
        plt.show()
        # fs = 44100
        # nperseg = int(0.02 * fs)
        # noverlap = int(0.01 * fs)
        # nfft = 1024
        # print(len(speedAll[i]))
        # f, t, Zxx = signal.stft(curSig, fs, nperseg=nperseg, noverlap=noverlap, nfft=nfft)
        # Zxx = np.abs(Zxx)
        # plt.pcolormesh(t, f, Zxx, cmap='YlOrRd')
        # plt.show()

    # speedRes = wiener(speedRes)
    # speedRes = WaveletTrans(speedRes, 0.5)
    speedRes = GaussSmooth(speedRes, 5, 3)

    timeShow = timeShow[:len(timeShow)-1]
    speedRes = speedRes[0:len(timeShow)]
    if count != None:
        plt.subplot(211+count)
        plt.plot(timeShow, speedRes)
        plt.xlabel('time(s)')
        plt.ylabel('speed(m/s)')
        plt.title(path)
        if count == 1:
            plt.show()
    else:
        plt.plot(timeShow, speedRes)
        plt.xlabel('time(s)')
        plt.ylabel('speed(m/s)')
        plt.title(path)
        plt.show()



def ShowGraph(dirPath, fc, lowTh, nfft, frameTime, aheadTime):

    pattern = os.path.join(dirPath, '*')
    count = 0
    for filePath in glob.glob(pattern):
        SingleFileDop(filePath, count)
        count = (count + 1) % 2
        # CalLipSig(
        #     filePath=filePath,
        #     fc=fc,
        #     lowTh=lowTh,
        #     nfft=nfft,
        #     frameTime=frameTime,
        #     aheadTime=aheadTime
        # )

dirPath = 'F:\\dataset\\mix_data10\\slow'
filePath = 'F:\\dataset\\mix_data10\\slow\\am_2320k.wav'
fc = 20000
lowTh = 100
nfft = 8192
frameTime = 0.15
aheadTime = 0.01
# ShowGraph(
#     dirPath=dirPath,
#     fc=fc,
#     lowTh=lowTh,
#     nfft=nfft,
#     frameTime=frameTime,
#     aheadTime=aheadTime
# )

SingleFileDop(filePath)
# ShowGraph('mixsig')
# GetCw('', 15550, 18350, 350)
